package com.lingjtx.ai.service.impl;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.lingjtx.ai.po.AiChatQueryPo;
import com.lingjtx.ai.po.ChatMsgFeedbackPo;
import com.lingjtx.ai.service.ChatService;
import com.lingjtx.common.core.constant.Const;
import com.lingjtx.common.core.constant.RedisConfig;
import com.lingjtx.common.core.util.StringUtil;
import com.lingjtx.common.redis.service.IRedisService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Service
public class ChatServiceImpl implements ChatService {

    @Autowired
    private IRedisService redisService;

    @Resource(name = "difyClient")
    private WebClient client;

    private final Map<String, Sinks.Many<ServerSentEvent<String>>> emitterPool = new ConcurrentHashMap<>();

    @Override
    public Flux<ServerSentEvent<String>> completions(AiChatQueryPo model, ServerHttpRequest request) {

        // 临时key：user + requestId，仅用于首次建立 sink 和返回 Flux
        String tempEmitterKey = model.getUser() + "_" + request.getId();

        // 创建新连接
        Sinks.Many<ServerSentEvent<String>> sink = Sinks.many().unicast().onBackpressureBuffer();

        client.post()
                .uri("/chat-messages")
                .accept(MediaType.TEXT_EVENT_STREAM)
                .bodyValue(JSON.toJSONString(model))
                .retrieve()
                .bodyToFlux(ServerSentEvent.class)
                .doOnNext(event -> {
                    if (event.data() != null) {
                        String llmModel = Const.LLMModel.deepseek_r1;
                        if (!Integer.valueOf(1).equals(model.getInputs().getR1())) {
                            llmModel = Const.LLMModel.deepseek_v3;
                        }
                        processDifyEvent(JSON.toJSONString(event.data()), model.getUser(), sink, llmModel);
                    }
                })
                .doOnError(error -> {
                    log.error("SSE 流出错: {}", error.getMessage(), error);
                    sink.tryEmitError(error);
                })
                .doOnComplete(() -> {
                    log.info("SSE 流结束: {}", tempEmitterKey);
                    sink.tryEmitComplete();
                })
                .subscribe();

        return sink.asFlux()
                .doOnCancel(() -> {
                    log.info("客户端主动关闭连接: {}", tempEmitterKey);
                    emitterPool.remove(tempEmitterKey); // 预防未获取 taskId 也能清除
                });
    }

    @Override
    public Mono<String> stop(String taskId, ChatMsgFeedbackPo model) {
        String emitterKey = model.getUser() + "_" + taskId;
        Sinks.Many<ServerSentEvent<String>> sink = emitterPool.remove(emitterKey);
        if (sink != null) {
            sink.tryEmitComplete();
        }

        return client.post()
                .uri(uriBuilder -> uriBuilder
                        .path("/chat-messages/{taskId}/stop")
                        .build(taskId))
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(model)
                .retrieve()
                .bodyToMono(String.class);
    }

    @Override
    public Mono<String> feedback(String message_id, ChatMsgFeedbackPo model) {
        return client.post()
                .uri(uriBuilder -> uriBuilder
                        .path("/messages/{message_id}/feedbacks")
                        .build(message_id))
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(model) // 自动序列化为JSON
                .retrieve()
                .bodyToMono(String.class);
    }

    private void processDifyEvent(String difyData, String user, Sinks.Many<ServerSentEvent<String>> sink, String model) {
        String taskId = null;
        try {
            // 解析 Dify 的 JSON 数据
            JSONObject difyJson = JSONObject.parseObject(difyData);
            String eventType = difyJson.getString("event");
            taskId = difyJson.getString("task_id");

            if (StringUtil.isNotEmpty(taskId)) {
                String emitterKey = user + "_" + taskId;
                emitterPool.putIfAbsent(emitterKey, sink);
            }

            Sinks.Many<ServerSentEvent<String>> currentSink = emitterPool.get(user + "_" + taskId);
            if (currentSink == null) {
                log.warn("Sink 不存在: user={}, taskId={}", user, taskId);
                return;
            }

            if ("workflow_started".equals(eventType)) { // 工作流开始
                redisService.set(RedisConfig.LLM_CHAT_TASK_ID + taskId, Const.LLMDeltaType.TEXT, RedisConfig.TIMEOUT_30_MIN);
            }

            // 处理消息事件
            if ("message".equals(eventType)) {
                String content = difyJson.getString("answer");
                String type = (String) redisService.get(RedisConfig.LLM_CHAT_TASK_ID + taskId);
                if (StringUtil.isNotEmpty(content)) {
                    if (content.contains("</ref>")) {
                        type = Const.LLMDeltaType.REF;
                    }
                    if (content.contains("<think>")) {
                        difyJson.put("answer", content.replace("<think>", ""));
                        type = Const.LLMDeltaType.THINKING;
                        redisService.set(RedisConfig.LLM_CHAT_TASK_ID + taskId, Const.LLMDeltaType.THINKING);
                    }
                    if (content.contains("<details style=\"color:gray;background-color: #f8f8f8;padding: 8px;border-radius: 4px;\" open> <summary> Thinking... </summary>")) {
                        difyJson.put("answer", content.replace("<details style=\"color:gray;background-color: #f8f8f8;padding: 8px;border-radius: 4px;\" open> <summary> Thinking... </summary>", ""));
                        type = Const.LLMDeltaType.THINKING;
                        redisService.set(RedisConfig.LLM_CHAT_TASK_ID + taskId, Const.LLMDeltaType.THINKING);
                    }

                    if (content.contains("</think>")) {
                        String[] str = content.split("</think>");
                        if (str.length > 1) {
                            difyJson.put("answer", str[0]);
                            type = Const.LLMDeltaType.THINKING;
                            redisService.set(RedisConfig.LLM_CHAT_TASK_ID + taskId, Const.LLMDeltaType.TEXT);
                            sink.tryEmitNext(ServerSentEvent.builder(convertToOpenAIChunk(difyJson, type, model)).build());

                            difyJson.put("answer", str[1]);
                            type = Const.LLMDeltaType.TEXT;
                        } else {
                            difyJson.put("answer", content.replace("</think>", ""));
                            type = Const.LLMDeltaType.THINKING;
                            redisService.set(RedisConfig.LLM_CHAT_TASK_ID + taskId, Const.LLMDeltaType.TEXT);
                        }
                    }
                    if (content.contains("</details>")) {
                        String[] str = content.split("</details>");
                        if (str.length > 1) {
                            difyJson.put("answer", str[0]);
                            type = Const.LLMDeltaType.THINKING;
                            redisService.set(RedisConfig.LLM_CHAT_TASK_ID + taskId, Const.LLMDeltaType.TEXT);
                            sink.tryEmitNext(ServerSentEvent.builder(convertToOpenAIChunk(difyJson, type, model)).build());

                            difyJson.put("answer", str[1]);
                            type = Const.LLMDeltaType.TEXT;
                        } else {
                            difyJson.put("answer", content.replace("</details>", ""));
                            type = Const.LLMDeltaType.THINKING;
                            redisService.set(RedisConfig.LLM_CHAT_TASK_ID + taskId, Const.LLMDeltaType.TEXT);
                        }
                    }
                    if (StringUtil.isEmpty(type)) {
                        type = Const.LLMDeltaType.TEXT;
                    }
                }
                String openAIData = convertToOpenAIChunk(difyJson, type, model);
                sink.tryEmitNext(ServerSentEvent.builder(openAIData).build());
            }
            // 处理结束事件
            else if ("workflow_finished".equals(eventType)) {
                // 本条对话结束，插入数据库操作
//                    String content = difyJson.getJSONObject("data").getJSONObject("outputs").getString("answer");
            } else if ("message_end".equals(eventType)) {
                sink.tryEmitNext(ServerSentEvent.builder(buildFinalResponse(difyJson, model)).build());
                sink.tryEmitNext(ServerSentEvent.builder("[DONE]").build());
                sink.tryEmitComplete();
                emitterPool.remove(user + "_" + taskId);
                redisService.delete(RedisConfig.LLM_CHAT_TASK_ID + taskId);
            }
        } catch (Exception e) {
            log.error("processDifyEvent 异常：{}", e.getMessage(), e);
            sink.tryEmitError(e);
            emitterPool.remove(user + "_" + taskId);
        }
    }

    private String convertToOpenAIChunk(JSONObject difyJson, String type, String model) {
        //{
        //    "event": "message",
        //    "conversation_id": "d278baaa-49f5-4660-9c5e-d563b9c147b5",
        //    "message_id": "61aa6a52-481c-43d4-b1cf-97b2188f638b",
        //    "created_at": 1742782962,
        //    "task_id": "17aca9c0-6394-4180-9509-6df463523d4c",
        //    "id": "61aa6a52-481c-43d4-b1cf-97b2188f638b",
        //    "answer": "<ref>",
        //    "from_variable_selector": [
        //        "1742778350648",
        //        "answer"
        //    ]
        //}

        //{
        //  "id": "930c60df-bf64-41c9-a88e-3ec75f81e00e",
        //  "choices": [
        //    {
        //      "finish_reason": "stop",
        //      "index": 0,
        //      "message": {
        //        "content": "Hello! How can I help you today?",
        //        "role": "assistant"
        //      }
        //    }
        //  ],
        //  "created": 1705651092,
        //  "model": "deepseek-chat",
        //  "object": "chat.completion",
        //  "usage": {
        //    "completion_tokens": 10,
        //    "prompt_tokens": 16,
        //    "total_tokens": 26
        //  }
        //}
        JSONObject chunk = new JSONObject();
        chunk.put("id", difyJson.getString("message_id")); // 使用 Dify 的 message_id
        chunk.put("conversation_id", difyJson.getString("conversation_id"));
        chunk.put("task_id", difyJson.getString("task_id"));
        chunk.put("object", "chat.completion.chunk");
        chunk.put("created", difyJson.getLong("created_at")); // 直接使用 Dify 的时间戳
        chunk.put("model", model); // 模型名称

        // 构造 choices 数组
        JSONArray choices = new JSONArray();
        JSONObject choice = new JSONObject();
        JSONObject delta = new JSONObject();
        if (Const.LLMDeltaType.THINKING.equals(type)) { // 兼容旧版本
            delta.put("reasoning_content", difyJson.getString("answer"));
        } else {
            delta.put("content", difyJson.getString("answer")); // 映射 Dify 的 answer 字段
        }
        delta.put("type", type);
        choice.put("delta", delta);
        choice.put("index", 0);
        choice.put("finish_reason", null); // 流式块中应为 null
        choices.add(choice);
        chunk.put("choices", choices);

        return chunk.toJSONString();
    }

    private String buildFinalResponse(JSONObject difyJson, String model) {
        //{
        //    "event": "message_end",
        //    "conversation_id": "b790f823-4a38-4785-8271-6bdcc6a14337",
        //    "message_id": "3c575ca3-32b4-4ab0-ac38-8293e6cf3318",
        //    "created_at": 1742808857,
        //    "task_id": "57537ecf-4d7c-4cc7-9bc2-66e996f1ca8d",
        //    "id": "3c575ca3-32b4-4ab0-ac38-8293e6cf3318",
        //    "metadata": {
        //        "usage": {
        //            "prompt_tokens": 1958,
        //            "prompt_unit_price": "4",
        //            "prompt_price_unit": "0.000001",
        //            "prompt_price": "0.007832",
        //            "completion_tokens": 1315,
        //            "completion_unit_price": "16",
        //            "completion_price_unit": "0.000001",
        //            "completion_price": "0.02104",
        //            "total_tokens": 3273,
        //            "total_price": "0.028872",
        //            "currency": "RMB",
        //            "latency": 111.29823345600016
        //        }
        //    },
        //    "files": []
        //}
        JSONObject finalResponse = new JSONObject();
        finalResponse.put("id", difyJson.getString("message_id"));
        finalResponse.put("object", "chat.completion");
        finalResponse.put("created", difyJson.getLong("created_at"));
        finalResponse.put("model", model);

        // 构造 choices
        JSONArray choices = new JSONArray();
        JSONObject choice = new JSONObject();
        JSONObject message = new JSONObject();
        message.put("content", difyJson.getString("answer"));
        message.put("role", "assistant");
        choice.put("message", message);
        choice.put("finish_reason", "stop");
        choices.add(choice);
        finalResponse.put("choices", choices);

        // 添加 usage
        JSONObject usage = new JSONObject();
        JSONObject difyUsage = difyJson.getJSONObject("metadata").getJSONObject("usage");
        usage.put("total_tokens", difyUsage.getIntValue("total_tokens"));
        usage.put("completion_tokens", difyUsage.getIntValue("completion_tokens"));
        usage.put("prompt_tokens", difyUsage.getIntValue("prompt_tokens"));
        finalResponse.put("usage", usage);

        return finalResponse.toJSONString();
    }
}